{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/chowgi/GenAI-Showcase/blob/main/mongodb_as_a_toolbox_for_llamaindex_agents.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af9cac0d"
      },
      "source": [
        "# MongoDB as a Toolbox for LlamaIndex Agents\n",
        "\n",
        "This notebook demonstrates how to leverage MongoDB Atlas as a \"toolbox\" for LlamaIndex agents. The application showcases the integration of MongoDB's capabilities, specifically its Vector Search feature, with LlamaIndex for building intelligent agents capable of performing various tasks by calling relevant tools stored and managed within MongoDB.\n",
        "\n",
        "**Key Features:**\n",
        "\n",
        "*   **MongoDB as a Tool Registry:** Instead of hardcoding tool definitions within the agent, this application stores tool metadata (name, description, parameters) directly in a MongoDB collection.\n",
        "*   **MongoDB Atlas Vector Search for Tool Discovery:** LlamaIndex uses the vector embeddings of tool descriptions stored in MongoDB to perform semantic searches based on user queries. This allows the agent to dynamically discover and select the most relevant tools for a given task.\n",
        "*   **LlamaIndex Agent with Function Calling:** The LlamaIndex agent is configured to use the retrieved tool definitions from MongoDB to enable function calling. This means the agent can understand the user's intent and execute the appropriate Python function (tool) stored in the application.\n",
        "*   **Data Storage in MongoDB:** Besides tool definitions, the application also uses separate MongoDB collections to store operational data like customer orders, return requests, and policy documents.\n",
        "*   **Integration with External Services:** The tools defined and managed in MongoDB can interact with external services (e.g., fetching real-time data, processing requests) or perform operations on the data stored within MongoDB itself (e.g., looking up order details, creating return requests).\n",
        "\n",
        "This approach provides a flexible and scalable way to manage and expand the agent's capabilities. New tools can be added to the MongoDB collection dynamically, and the agent can discover and utilize them without requiring code changes to the agent itself."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bc3f1647"
      },
      "source": [
        "# Environment Setup and Configuration\n",
        "\n",
        "This section covers the installation of necessary libraries, setting up API keys, and configuring the database connection to MongoDB Atlas.\n",
        "\n",
        "### Install required libraries\n",
        "\n",
        "This cell installs the necessary Python libraries using `uv pip install`. These libraries include:\n",
        "- `pymongo`: A Python driver for MongoDB.\n",
        "- `llama-index-core`: The core LlamaIndex library.\n",
        "- `llama-index-llms-openai`: LlamaIndex integration with OpenAI LLMs.\n",
        "- `llama-index-embeddings-voyageai`: LlamaIndex integration with VoyageAI embeddings.\n",
        "- `llama-index-vector-stores-mongodb`: LlamaIndex integration with MongoDB Atlas Vector Search.\n",
        "- `llama-index-readers-file`: LlamaIndex file readers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6s3dlQKnRkFL"
      },
      "outputs": [],
      "source": [
        "!uv pip install pymongo llama-index-core llama-index-llms-openai llama-index-embeddings-voyageai llama-index-vector-stores-mongodb llama-index-readers-file"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7f9a1c5a"
      },
      "source": [
        "### Get and store API keys\n",
        "\n",
        "Get and store API keys\n",
        "This cell retrieves API keys for OpenAI, MongoDB, and VoyageAI from Google Colab's user data secrets and sets them as environment variables.\n",
        "\n",
        "Please obtain your own API keys for OpenAI, MongoDB Atlas, and VoyageAI.\n",
        "\n",
        "OpenAI: You can get an API key from the OpenAI website.\n",
        "MongoDB Atlas: Get your connection string from your MongoDB Atlas cluster.\n",
        "VoyageAI: Obtain an API key from the VoyageAI website.\n",
        "Once you have your keys, add them to Google Colab's user data secrets by clicking on the \"🔑\" icon in the left sidebar. Name the secrets OPENAI_API_KEY, MONGODB_URI, and VOYAGE_API_KEY respectively.\n",
        "\n",
        "It also defines the GPT model to be used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mj2FLQpkUKcl"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from google.colab import userdata\n",
        "\n",
        "OPENAI_API_KEY = userdata.get(\"OPENAI_API_KEY\")\n",
        "os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY\n",
        "\n",
        "MONGO_URI = userdata.get(\"MONGODB_URI\")\n",
        "os.environ[\"MONGO_URI\"] = MONGO_URI\n",
        "\n",
        "VOYAGE_API_KEY = userdata.get(\"VOYAGE_API_KEY\")\n",
        "os.environ[\"VOYAGE_API_KEY\"] = VOYAGE_API_KEY\n",
        "\n",
        "GPT_MODEL = \"gpt-4o\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f45d745a"
      },
      "source": [
        "### Setup the database\n",
        "\n",
        "This cell establishes a connection to the MongoDB Atlas database using the provided URI. It then defines the database name and the names of the collections that will be used in this notebook for storing tools, orders, returns, and policies. Finally, it creates client objects for each of these collections."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM0G-evHCMUv"
      },
      "outputs": [],
      "source": [
        "import pymongo\n",
        "\n",
        "# Get MongoClient\n",
        "mongo_client = pymongo.MongoClient(MONGO_URI, appname=\"showcase.tools.mongodb_toolbox\")\n",
        "\n",
        "# Set the DB name\n",
        "db_name = \"retail_agent_demo\"\n",
        "\n",
        "# Set the database client\n",
        "db = mongo_client[db_name]\n",
        "\n",
        "# Set the required collection names\n",
        "tools_collection_name = \"tools\"\n",
        "orders_collection_name = \"orders\"\n",
        "returns_collection_name = \"returns\"\n",
        "policies_collection_name = \"policies\"\n",
        "\n",
        "tools_collection = db[tools_collection_name]\n",
        "orders_collection = db[orders_collection_name]\n",
        "returns_collection = db[returns_collection_name]\n",
        "policies_collection = db[policies_collection_name]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2edCDr1nZ_my"
      },
      "source": [
        "# Loading Demo Data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6f3012b"
      },
      "source": [
        "## Download and store policy documents into MongoDB Atlas vector store\n",
        "\n",
        "This cell downloads policy documents and stores them in a MongoDB Atlas vector store. It initializes a vector store, checks if the collection is empty, downloads PDF documents, loads them, adds metadata, initializes embedding and node parsing, parses documents into nodes, creates a storage context, creates a vector index, and ingests the documents."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Klh_JQlDRl8Y"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import requests\n",
        "from llama_index.core import StorageContext, VectorStoreIndex\n",
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "from llama_index.embeddings.voyageai import VoyageEmbedding\n",
        "\n",
        "# Import PDFReader\n",
        "from llama_index.readers.file import PDFReader\n",
        "from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch\n",
        "\n",
        "# Set up vector store for the policies collection\n",
        "policy_vector_store = MongoDBAtlasVectorSearch(\n",
        "    mongo_client,\n",
        "    db_name=db_name,\n",
        "    collection_name=\"policies\",\n",
        "    vector_index_name=\"vector_index\",  # Assuming a vector index named 'vector_index' exists\n",
        ")\n",
        "\n",
        "# Check if the policies collection is empty\n",
        "policies_count = policy_vector_store.collection.count_documents({})\n",
        "if policies_count > 0:\n",
        "    print(\n",
        "        f\"Policies collection is not empty. Skipping document import. Total documents: {policies_count}\"\n",
        "    )\n",
        "else:\n",
        "    print(\"Policies collection is empty. Starting document import.\")\n",
        "\n",
        "    # Define the list of document URLs\n",
        "    document_urls = [\n",
        "        \"https://mongodb-llamaindex-demos.s3.us-west-1.amazonaws.com/privacy_policy.pdf\",\n",
        "        \"https://mongodb-llamaindex-demos.s3.us-west-1.amazonaws.com/return_policy.pdf\",\n",
        "        \"https://mongodb-llamaindex-demos.s3.us-west-1.amazonaws.com/shipping_policy.pdf\",\n",
        "        \"https://mongodb-llamaindex-demos.s3.us-west-1.amazonaws.com/terms_of_service.pdf\",\n",
        "        \"https://mongodb-llamaindex-demos.s3.us-west-1.amazonaws.com/warranty_policy.pdf\",\n",
        "    ]\n",
        "\n",
        "    # Create a temporary directory to store the downloaded files\n",
        "    temp_dir = \"temp_policy_docs\"\n",
        "    os.makedirs(temp_dir, exist_ok=True)\n",
        "\n",
        "    # Download each file to the temporary directory\n",
        "    local_files = []\n",
        "    for url in document_urls:\n",
        "        file_name = os.path.join(temp_dir, url.split(\"/\")[-1])\n",
        "        try:\n",
        "            response = requests.get(url)\n",
        "            response.raise_for_status()  # Raise an HTTPError for bad responses\n",
        "            with open(file_name, \"wb\") as f:\n",
        "                f.write(response.content)\n",
        "            local_files.append(file_name)\n",
        "            print(f\"Downloaded {url} to {file_name}\")\n",
        "        except requests.exceptions.RequestException as e:\n",
        "            print(f\"Error downloading {url}: {e}\")\n",
        "\n",
        "    # Use PDFReader to load each PDF file from the temporary directory\n",
        "    documents = []\n",
        "    for file_path in local_files:\n",
        "        try:\n",
        "            loader = PDFReader()\n",
        "            docs = loader.load_data(file=file_path)\n",
        "            documents.extend(docs)\n",
        "            print(f\"Loaded {file_path}\")\n",
        "        except Exception as e:\n",
        "            print(f\"Error loading {file_path} with PDFReader: {e}\")\n",
        "\n",
        "    # Add metadata to documents (optional, but can be useful)\n",
        "    for i, doc in enumerate(documents):\n",
        "        doc.metadata.update(\n",
        "            {\n",
        "                \"document_type\": \"policy\",\n",
        "                \"document_index\": i,\n",
        "                \"file_name\": os.path.basename(doc.metadata.get(\"file_path\", \"unknown\")),\n",
        "            }\n",
        "        )\n",
        "\n",
        "    print(f\"Loaded {len(documents)} documents from directory\")\n",
        "\n",
        "    # Initialize embedding model\n",
        "    embed_model = VoyageEmbedding(model_name=\"voyage-3.5-lite\", api_key=VOYAGE_API_KEY)\n",
        "\n",
        "    # Initialize node parser for chunking\n",
        "    node_parser = SentenceSplitter(chunk_size=2024, chunk_overlap=200)\n",
        "\n",
        "    # Parse documents into nodes (chunks)\n",
        "    nodes = node_parser.get_nodes_from_documents(documents)\n",
        "    print(f\"Created {len(nodes)} text chunks\")\n",
        "\n",
        "    # Create storage context with MongoDB vector store\n",
        "    storage_context = StorageContext.from_defaults(vector_store=policy_vector_store)\n",
        "\n",
        "    # Create vector index and ingest documents into MongoDB\n",
        "    # This step automatically adds nodes to the vector store when storage_context is provided\n",
        "    index = VectorStoreIndex(\n",
        "        nodes=nodes,\n",
        "        storage_context=storage_context,\n",
        "        embed_model=embed_model,\n",
        "        show_progress=True,\n",
        "    )\n",
        "\n",
        "    print(\"Successfully ingested all PDF documents into MongoDB 'policies' collection\")\n",
        "\n",
        "    # Display collection stats using the vector store's collection object\n",
        "    policies_count = policy_vector_store.collection.count_documents({})\n",
        "    print(f\"Total documents in 'policies' collection: {policies_count}\")\n",
        "\n",
        "    # Optional: Clean up the temporary directory\n",
        "    # import shutil\n",
        "    # shutil.rmtree(temp_dir)\n",
        "    # print(f\"Cleaned up temporary directory: {temp_dir}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3f65d6d"
      },
      "source": [
        "## Create and Store Dummy Order Data\n",
        "\n",
        "This cell generates a list of fake order data with details like order ID, date, status, total amount, shipping address, payment method, and items. It then checks if the `orders` collection in MongoDB is empty. If it is, the fake order data is inserted into the `orders` collection. This is done to populate the database with sample data for testing and demonstrating the order lookup functionality later in the notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "84044d5d"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "# Check if the collection is empty\n",
        "if orders_collection.count_documents({}) == 0:\n",
        "    # Define some fake order data\n",
        "    fake_orders = [\n",
        "        {\n",
        "            \"order_id\": 101,\n",
        "            \"order_date\": datetime(2023, 10, 26, 10, 0, 0),\n",
        "            \"status\": \"Shipped\",\n",
        "            \"total_amount\": 150.75,\n",
        "            \"shipping_address\": \"123 Main St, Anytown, CA 91234\",\n",
        "            \"payment_method\": \"Credit Card\",\n",
        "            \"items\": [\n",
        "                {\"name\": \"Laptop\", \"price\": 1200.00},\n",
        "                {\"name\": \"Mouse\", \"price\": 25.75},\n",
        "            ],\n",
        "        },\n",
        "        {\n",
        "            \"order_id\": 102,\n",
        "            \"order_date\": datetime(2023, 10, 25, 14, 30, 0),\n",
        "            \"status\": \"Processing\",\n",
        "            \"total_amount\": 55.00,\n",
        "            \"shipping_address\": \"456 Oak Ave, Somewhere, NY 54321\",\n",
        "            \"payment_method\": \"PayPal\",\n",
        "            \"items\": [\n",
        "                {\"name\": \"Keyboard\", \"price\": 75.00},\n",
        "            ],\n",
        "        },\n",
        "        {\n",
        "            \"order_id\": 103,\n",
        "            \"order_date\": datetime(2023, 10, 25, 14, 30, 0),\n",
        "            \"status\": \"Processing\",\n",
        "            \"total_amount\": 35.00,\n",
        "            \"shipping_address\": \"789 Pine Rd, Elsewhere, TX 67890\",\n",
        "            \"payment_method\": \"Debit Card\",\n",
        "            \"items\": [\n",
        "                {\"name\": \"Monitor\", \"price\": 250.00},\n",
        "            ],\n",
        "        },\n",
        "    ]\n",
        "\n",
        "    # Insert the fake orders into the collection\n",
        "    orders_collection.insert_many(fake_orders)\n",
        "    print(f\"Inserted {len(fake_orders)} fake orders.\")\n",
        "else:\n",
        "    print(\"Orders collection is not empty. Skipping insertion of fake orders.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6eMnUDZFaTZB"
      },
      "source": [
        "# Application Setup and Configuration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "218f57e9"
      },
      "source": [
        "## Define MongoDB Tool Decorator\n",
        "\n",
        "This cell defines the `mongodb_toolbox` decorator. This decorator is used to register functions as tools that can be discovered and used by the LlamaIndex agent. It also handles generating embeddings for the tool descriptions and storing them in the MongoDB 'tools' collection for vector search."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EJe1J0QOB_AB"
      },
      "outputs": [],
      "source": [
        "import inspect\n",
        "from functools import wraps\n",
        "from typing import get_type_hints\n",
        "\n",
        "from llama_index.core import Document, StorageContext, VectorStoreIndex\n",
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "from llama_index.embeddings.voyageai import VoyageEmbedding\n",
        "from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch\n",
        "\n",
        "# Initialize vector store\n",
        "vector_store = MongoDBAtlasVectorSearch(\n",
        "    mongo_client,\n",
        "    db_name=db_name,\n",
        "    collection_name=\"tools\",\n",
        "    vector_index_name=\"vector_index\",\n",
        ")\n",
        "\n",
        "# Initialize VoyageAIEmbedding\n",
        "voyage_embed_model = VoyageEmbedding(\n",
        "    model_name=\"voyage-3.5-lite\",\n",
        ")\n",
        "\n",
        "# Create a registry for decorated tools\n",
        "decorated_tools_registry = {}\n",
        "\n",
        "\n",
        "def get_embedding(text):\n",
        "    text = text.replace(\"\\n\", \" \")\n",
        "    return voyage_embed_model.get_text_embedding(text)\n",
        "\n",
        "\n",
        "def mongodb_toolbox(vector_store=None):\n",
        "    def decorator(func):\n",
        "        @wraps(func)\n",
        "        def wrapper(*args, **kwargs):\n",
        "            return func(*args, **kwargs)\n",
        "\n",
        "        # Generate tool definition\n",
        "        signature = inspect.signature(func)\n",
        "        docstring = inspect.getdoc(func) or \"\"\n",
        "        type_hints = get_type_hints(func)\n",
        "\n",
        "        tool_def = {\n",
        "            \"name\": func.__name__,\n",
        "            \"description\": docstring.strip(),\n",
        "            \"parameters\": {\"type\": \"object\", \"properties\": {}, \"required\": []},\n",
        "        }\n",
        "\n",
        "        for param_name, param in signature.parameters.items():\n",
        "            if (\n",
        "                param.kind == inspect.Parameter.VAR_POSITIONAL\n",
        "                or param.kind == inspect.Parameter.VAR_KEYWORD\n",
        "            ):\n",
        "                continue\n",
        "\n",
        "            param_type = type_hints.get(param_name, type(None))\n",
        "            json_type = \"string\"  # Default to string\n",
        "            if param_type in (int, float):\n",
        "                json_type = \"number\"\n",
        "            elif param_type is bool:\n",
        "                json_type = \"boolean\"\n",
        "\n",
        "            tool_def[\"parameters\"][\"properties\"][param_name] = {\n",
        "                \"type\": json_type,\n",
        "                \"description\": f\"Parameter {param_name}\",\n",
        "            }\n",
        "\n",
        "            if param.default == inspect.Parameter.empty:\n",
        "                tool_def[\"parameters\"][\"required\"].append(param_name)\n",
        "\n",
        "        tool_def[\"parameters\"][\"additionalProperties\"] = False\n",
        "\n",
        "        # Create Document for vector storage with embedding\n",
        "        document = Document(text=tool_def[\"description\"], metadata=tool_def)\n",
        "\n",
        "        # Generate and set the embedding\n",
        "        if vector_store and tool_def[\"description\"]:\n",
        "            embedding = voyage_embed_model.get_text_embedding(tool_def[\"description\"])\n",
        "            document.embedding = embedding\n",
        "\n",
        "        # Add to vector store only if a document with the same name does not exist\n",
        "        if vector_store:\n",
        "            existing_doc = vector_store.collection.find_one(\n",
        "                {\"metadata.name\": tool_def[\"name\"]}\n",
        "            )\n",
        "            if not existing_doc:\n",
        "                vector_store.add([document])\n",
        "            else:\n",
        "                print(\n",
        "                    f\"Document for tool '{tool_def['name']}' already exists. Skipping insertion.\"\n",
        "                )\n",
        "\n",
        "        # Register the decorated function\n",
        "        decorated_tools_registry[func.__name__] = func\n",
        "\n",
        "        return wrapper\n",
        "\n",
        "    return decorator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c46ecd38"
      },
      "source": [
        "## Setup indexes\n",
        "\n",
        "This cell checks for and creates vector search indexes on the specified MongoDB collections if they don't already exist. These indexes are crucial for performing efficient vector searches on the data stored in these collections."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bkLgKvTozKdn"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "# Require vector index list\n",
        "required_indexs = [\n",
        "    orders_collection_name,\n",
        "    tools_collection_name,\n",
        "    returns_collection_name,\n",
        "    policies_collection_name,\n",
        "]\n",
        "\n",
        "# Flag to track if any index was created\n",
        "index_created = False\n",
        "\n",
        "for collection_name in required_indexs:\n",
        "    print(f\"Checking and creating index for collection: {collection_name}\")\n",
        "    # Set up vector store for the current collection\n",
        "    current_vector_store = MongoDBAtlasVectorSearch(\n",
        "        mongo_client,\n",
        "        db_name=db_name,\n",
        "        collection_name=collection_name,\n",
        "        vector_index_name=\"vector_index\",\n",
        "    )\n",
        "\n",
        "    # Check if vector index exists\n",
        "    try:\n",
        "        search_indexes = list(current_vector_store.collection.list_search_indexes())\n",
        "        index_exists = any(\n",
        "            index.get(\"name\") == \"vector_index\" for index in search_indexes\n",
        "        )\n",
        "    except Exception as e:\n",
        "        print(f\"Could not check search indexes for {collection_name}: {e}\")\n",
        "        index_exists = False\n",
        "\n",
        "    if not index_exists:\n",
        "        # Index does not exist, create it\n",
        "        current_vector_store.create_vector_search_index(\n",
        "            dimensions=1024, path=\"embedding\", similarity=\"cosine\"\n",
        "        )\n",
        "        print(f\"Vector search index created successfully for {collection_name}.\")\n",
        "        index_created = True  # Set flag if an index was created\n",
        "    else:\n",
        "        # Index exists, skip creation\n",
        "        print(\n",
        "            f\"Vector search index already exists for {collection_name}. Skipping creation.\"\n",
        "        )\n",
        "\n",
        "# Add a single 20-second pause after checking all collections, only if an index was created\n",
        "if index_created:\n",
        "    print(\"Pausing for 20 seconds to allow index builds...\")\n",
        "    time.sleep(20)\n",
        "    print(\"Resuming after pause.\")\n",
        "else:\n",
        "    print(\"No new indexes were created. Skipping pause.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "63b48ebf"
      },
      "source": [
        "## Define Vector Search Function\n",
        "\n",
        "This cell defines the `vector_search_tools` function, which performs a vector search on a given LlamaIndex vector store based on a user query. It uses the specified vector store and embedding model to find the most relevant documents (in this case, tool definitions) and returns a list of their metadata."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NmFa7RrZTKyk"
      },
      "outputs": [],
      "source": [
        "def vector_search_tools(user_query, vector_store, top_k=3):\n",
        "    \"\"\"\n",
        "    Perform a vector search using LlamaIndex vector store.\n",
        "\n",
        "    Args:\n",
        "        user_query (str): The user's query string.\n",
        "        vector_store: The LlamaIndex vector store instance.\n",
        "        top_k (int): Number of top results to return.\n",
        "\n",
        "    Returns:\n",
        "        list: A list of matching tool definitions.\n",
        "    \"\"\"\n",
        "    # Create index from vector store\n",
        "    index = VectorStoreIndex.from_vector_store(\n",
        "        vector_store,\n",
        "        embed_model=voyage_embed_model,\n",
        "    )\n",
        "\n",
        "    # Create query engine\n",
        "    query_engine = index.as_query_engine(similarity_top_k=top_k)\n",
        "\n",
        "    # Perform query\n",
        "    response = query_engine.query(user_query)\n",
        "\n",
        "    # Extract tool definitions from source nodes\n",
        "    tools_data = []\n",
        "    for node in response.source_nodes:\n",
        "        tool_metadata = node.node.metadata\n",
        "        tools_data.append(tool_metadata)\n",
        "\n",
        "    return tools_data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd8f99b6"
      },
      "source": [
        "## Define MongoDB Tools\n",
        "\n",
        "This cell defines several Python functions that will serve as tools for the LlamaIndex agent. Each function is decorated with the `@mongodb_toolbox` decorator, which registers the function and stores its definition and embedding in the 'tools' collection in MongoDB. These tools include functions for shouting, getting weather, getting stock price, getting current time, looking up orders, responding in Spanish, checking return policy, and creating a return request."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Uvko9J-0SCn4"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "from datetime import datetime\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def get_current_time(timezone: str = \"UTC\") -> str:\n",
        "    \"\"\"\n",
        "    Get the current time for a specified timezone.\n",
        "    Use this when a user asks about the current time in a specific timezone.\n",
        "\n",
        "    :param timezone: The timezone to get the current time for. Defaults to 'UTC'.\n",
        "    :return: A string with the current time in the specified timezone.\n",
        "    \"\"\"\n",
        "    current_time = datetime.utcnow().strftime(\"%H:%M:%S\")\n",
        "    return f\"The current time in {timezone} is {current_time}.\"\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def lookup_order_number(order_id: int) -> str:\n",
        "    \"\"\"\n",
        "    Lookup the details of a specific order number using its order ID.\n",
        "    Use this when a user asks for information about a particular order.\n",
        "\n",
        "    :param order_id: The unique identifier of the order to look up.\n",
        "    :return: A string containing the order details or a message if the order is not found.\n",
        "    \"\"\"\n",
        "    # Get the orders collection (assuming 'db' is accessible from here)\n",
        "    orders_collection = db[\"orders\"]\n",
        "\n",
        "    # Find the order by order_id\n",
        "    order = orders_collection.find_one(\n",
        "        {\"order_id\": order_id}\n",
        "    )  # Note: Sample data uses 'order_id'\n",
        "\n",
        "    if order:\n",
        "        # Format the order details into a readable string\n",
        "        order_details = f\"Order ID: {order.get('order_id')}\\n\"\n",
        "        order_details += (\n",
        "            f\"Order Date: {order.get('order_date').strftime('%Y-%m-%d %H:%M:%S')}\\n\"\n",
        "        )\n",
        "        order_details += f\"Status: {order.get('status')}\\n\"\n",
        "        order_details += f\"Total Amount: ${order.get('total_amount')}\\n\"\n",
        "        order_details += f\"Shipping Address: {order.get('shipping_address')}\\n\"\n",
        "        order_details += f\"Payment Method: {order.get('payment_method')}\\n\"\n",
        "        order_details += \"Items:\\n\"\n",
        "        for item in order.get(\"items\", []):\n",
        "            order_details += f\"- {item.get('name')}: ${item.get('price')}\\n\"\n",
        "\n",
        "        return order_details\n",
        "    else:\n",
        "        return f\"Order with ID {order_id} not found.\"\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def return_policy(return_request_description: str) -> str:\n",
        "    \"\"\"\n",
        "    Performs search on the policies collection to determine if a user's\n",
        "    return request aligns with the company's return policy, warranty policy,\n",
        "    or terms of service. Use this tool when a user asks about returning an item\n",
        "    and you need to check the relevant company policies.\n",
        "\n",
        "    Args:\n",
        "        return_request_description (str): A detailed description of the user's\n",
        "                                          return request, including reasons\n",
        "                                          for return, item condition, and any\n",
        "                                          relevant order information.\n",
        "\n",
        "    Returns:\n",
        "        str: A string containing relevant policy information found through\n",
        "             vector search that can help determine if the return request\n",
        "             meets the company's policy.\n",
        "    \"\"\"\n",
        "    # Create index from the policy vector store\n",
        "    policy_index = VectorStoreIndex.from_vector_store(\n",
        "        policy_vector_store,\n",
        "        embed_model=voyage_embed_model,\n",
        "    )\n",
        "\n",
        "    # Create query engine for the policy index\n",
        "    policy_query_engine = policy_index.as_query_engine(\n",
        "        similarity_top_k=3\n",
        "    )  # Adjust top_k as needed\n",
        "\n",
        "    # Perform query on the policies collection with the user's return request\n",
        "    response = policy_query_engine.query(return_request_description)\n",
        "\n",
        "    # Return the response text from the query engine\n",
        "    return str(response)\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def create_return_request(order_id: int, reason: str) -> str:\n",
        "    \"\"\"\n",
        "    Creates a return request entry for a given order ID with the specified reason.\n",
        "    Use this when a user wants to initiate a return for an item from a specific order.\n",
        "\n",
        "    :param order_id: The unique identifier of the order for which the return is requested.\n",
        "    :param reason: The reason for the return.\n",
        "    :return: A string confirming the return creation or indicating if the order was not found.\n",
        "    \"\"\"\n",
        "    # Get the orders and returns collections\n",
        "    orders_collection = db[\"orders\"]\n",
        "    returns_collection = db[\"returns\"]\n",
        "\n",
        "    # Find the order by order_id\n",
        "    order = orders_collection.find_one({\"order_id\": order_id})\n",
        "\n",
        "    if order:\n",
        "        # Create a return document\n",
        "        return_data = {\n",
        "            \"return_id\": returns_collection.count_documents({})\n",
        "            + 1,  # Simple auto-incrementing ID\n",
        "            \"order_id\": order_id,\n",
        "            \"return_date\": datetime.utcnow(),\n",
        "            \"reason\": reason,\n",
        "            \"status\": \"Pending\",  # Initial status\n",
        "            \"items\": order.get(\"items\", []),  # Include items from the original order\n",
        "        }\n",
        "\n",
        "        # Insert the return document into the returns collection\n",
        "        returns_collection.insert_one(return_data)\n",
        "\n",
        "        return f\"Return request created successfully for order {order_id} with reason: {reason}.\"\n",
        "    else:\n",
        "        return f\"Order with ID {order_id} not found. Could not create return.\"\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def greet_user(name: str) -> str:\n",
        "    \"\"\"\n",
        "    Greets the user by name.\n",
        "    Use this when a user provides their name and you want to greet them.\n",
        "\n",
        "    :param name: The name of the user.\n",
        "    :return: A greeting message.\n",
        "    \"\"\"\n",
        "    return f\"Hello, {name}! Nice to meet you.\"\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def calculate_square_root(number: float) -> str:\n",
        "    \"\"\"\n",
        "    Calculates the square root of a given number.\n",
        "    Use this when a user asks for the square root of a number.\n",
        "\n",
        "    :param number: The number for which to calculate the square root.\n",
        "    :return: A string with the square root result.\n",
        "    \"\"\"\n",
        "    if number < 0:\n",
        "        return \"Cannot calculate the square root of a negative number.\"\n",
        "    return f\"The square root of {number} is {number**0.5}.\"\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def repeat_phrase(phrase: str, times: int = 1) -> str:\n",
        "    \"\"\"\n",
        "    Repeats a given phrase a specified number of times.\n",
        "    Use this when a user asks you to repeat something.\n",
        "\n",
        "    :param phrase: The phrase to repeat.\n",
        "    :param times: The number of times to repeat the phrase. Defaults to 1.\n",
        "    :return: A string with the repeated phrase.\n",
        "    \"\"\"\n",
        "    if times <= 0:\n",
        "        return \"Please specify a positive number of times to repeat.\"\n",
        "    return (phrase + \" \") * times\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def roll_dice(number_of_dice: int = 1, sides: int = 6) -> str:\n",
        "    \"\"\"\n",
        "    Rolls a specified number of dice with a given number of sides and returns the results.\n",
        "    Use this when a user asks to roll dice.\n",
        "\n",
        "    :param number_of_dice: The number of dice to roll. Defaults to 1.\n",
        "    :param sides: The number of sides on each die. Defaults to 6.\n",
        "    :return: A string showing the result of each roll and the total.\n",
        "    \"\"\"\n",
        "    if number_of_dice <= 0 or sides <= 0:\n",
        "        return \"Please specify a positive number of dice and sides.\"\n",
        "    rolls = [random.randint(1, sides) for _ in range(number_of_dice)]\n",
        "    total = sum(rolls)\n",
        "    return f\"You rolled {number_of_dice} dice with {sides} sides each. Results: {rolls}. Total: {total}.\"\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def flip_coin(number_of_flips: int = 1) -> str:\n",
        "    \"\"\"\n",
        "    Flips a coin a specified number of times and returns the results.\n",
        "    Use this when a user asks to flip a coin.\n",
        "\n",
        "    :param number_of_flips: The number of times to flip the coin. Defaults to 1.\n",
        "    :return: A string showing the result of each flip.\n",
        "    \"\"\"\n",
        "    if number_of_flips <= 0:\n",
        "        return \"Please specify a positive number of flips.\"\n",
        "    results = [random.choice([\"Heads\", \"Tails\"]) for _ in range(number_of_flips)]\n",
        "    return f\"You flipped the coin {number_of_flips} times. Results: {results}.\"\n",
        "\n",
        "\n",
        "@mongodb_toolbox(vector_store=vector_store)\n",
        "def generate_random_password(length: int = 12) -> str:\n",
        "    \"\"\"\n",
        "    Generates a random password of a specified length.\n",
        "    Use this when a user asks for a random password.\n",
        "\n",
        "    :param length: The desired length of the password. Defaults to 12.\n",
        "    :return: A randomly generated password string.\n",
        "    \"\"\"\n",
        "    if length <= 0:\n",
        "        return \"Please specify a positive password length.\"\n",
        "    import string\n",
        "\n",
        "    characters = string.ascii_letters + string.digits + string.punctuation\n",
        "    password = \"\".join(random.choice(characters) for i in range(length))\n",
        "    return f\"Here is a random password: {password}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6f363495"
      },
      "source": [
        "## Populate Tools Function\n",
        "\n",
        "This cell defines the `populate_tools` function. This function takes the results from a vector search (which are tool definitions) and converts them into a list of LlamaIndex `FunctionTool` objects. It looks up the actual function object in the `decorated_tools_registry` based on the tool name found in the search results and creates a `FunctionTool` with the corresponding function and its description."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uurRWM6_TpUn"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.tools import FunctionTool\n",
        "\n",
        "# Access the registry created in the mongodb_toolbox decorator definition cell\n",
        "from __main__ import decorated_tools_registry\n",
        "\n",
        "\n",
        "def populate_tools(search_results):\n",
        "    \"\"\"\n",
        "    Populate the tools array based on the results from the vector search,\n",
        "    returning LlamaIndex FunctionTool objects.\n",
        "\n",
        "    Args:\n",
        "    search_results (list): The list of documents returned from the vector search.\n",
        "\n",
        "    Returns:\n",
        "    list: A list of LlamaIndex FunctionTool objects.\n",
        "    \"\"\"\n",
        "    tools = []\n",
        "\n",
        "    for result in search_results:\n",
        "        tool_name = result[\"name\"]\n",
        "        # Look up the function object in the registry\n",
        "        function_obj = decorated_tools_registry.get(tool_name)\n",
        "\n",
        "        if function_obj:\n",
        "            # Create a FunctionTool object from the actual function and its description\n",
        "            description = result.get(\"description\", function_obj.__doc__ or \"\")\n",
        "            tools.append(\n",
        "                FunctionTool.from_defaults(fn=function_obj, description=description)\n",
        "            )\n",
        "    return tools"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ksZ7zMFbWWl"
      },
      "source": [
        "# Running the Agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dc0262dc"
      },
      "source": [
        "## Test Tool Retrieval\n",
        "\n",
        "This cell demonstrates how to use the `vector_search_tools` function to find relevant tools based on a user query and then uses the `populate_tools` function to convert the search results into LlamaIndex `FunctionTool` objects. Finally, it prints the names of the retrieved tools to verify the process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PhOBBC_CSY5E"
      },
      "outputs": [],
      "source": [
        "import pprint\n",
        "\n",
        "user_query = (\n",
        "    \"I want to return a damaged laptop from order 101. What is the return policy??\"\n",
        ")\n",
        "\n",
        "tools_related_to_user_query = vector_search_tools(user_query, vector_store)\n",
        "\n",
        "# populate_tools now returns FunctionTool objects\n",
        "tools = populate_tools(tools_related_to_user_query)\n",
        "\n",
        "\n",
        "# Iterate through the list of FunctionTool objects and print their names\n",
        "tool_names = [tool.metadata.name for tool in tools]\n",
        "print(\n",
        "    \"Selected tools from the toolbox based on the similarity to the users intention -\"\n",
        ")\n",
        "pprint.pprint(tool_names)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3U5RoCzI6MTS"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.agent.workflow import FunctionAgent\n",
        "from llama_index.core.memory import Memory\n",
        "from llama_index.llms.openai import OpenAI\n",
        "\n",
        "# Setup The agent\n",
        "llm = OpenAI(model=GPT_MODEL)\n",
        "\n",
        "tools = populate_tools(tools_related_to_user_query)\n",
        "\n",
        "memory = Memory.from_defaults(session_id=\"my_session\", token_limit=40000)\n",
        "\n",
        "agent = FunctionAgent(llm=llm, tools=tools)\n",
        "\n",
        "# Get the answer\n",
        "response = await agent.run(\"How much did I pay for order 101\", memory=memory)\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T6EsN6nM9CvI"
      },
      "outputs": [],
      "source": [
        "response = await agent.run(\n",
        "    \"I want to return a damaged laptop from order 101. What is the return policy??\",\n",
        "    memory=memory,\n",
        ")\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6A8aMH9WPJi2"
      },
      "outputs": [],
      "source": [
        "response = await agent.run(\"Yes please\", memory=memory)\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8c69567f"
      },
      "source": [
        "### Get and store API keys\n",
        "\n",
        "This cell retrieves API keys for OpenAI, MongoDB, and VoyageAI from Google Colab's user data secrets and sets them as environment variables.\n",
        "\n",
        "**Please obtain your own API keys for OpenAI, MongoDB Atlas, and VoyageAI.**\n",
        "\n",
        "*   **OpenAI:** You can get an API key from the [OpenAI website](https://platform.openai.com/).\n",
        "*   **MongoDB Atlas:** Get your connection string from your MongoDB Atlas cluster.\n",
        "*   **VoyageAI:** Obtain an API key from the [VoyageAI website](https://voyageai.com/).\n",
        "\n",
        "Once you have your keys, add them to Google Colab's user data secrets by clicking on the \"🔑\" icon in the left sidebar. Name the secrets `OPENAI_API_KEY`, `MONGODB_URI`, and `VOYAGE_API_KEY` respectively.\n",
        "\n",
        "It also defines the GPT model to be used."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "state": {}
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
